#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：TU2.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/10/23 15:31 
'''
# !/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：algorithm_comparison.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/10/18 0:01 
'''
import os
import json
import random
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import time
from datetime import datetime
from typing import List, Tuple, Dict, Optional

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False


class GasSchedulingInstanceGenerator:
    """气体调度算例生成器与加载器"""

    @staticmethod
    def load_instance(file_path: str) -> Dict:
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"算例文件不存在: {file_path}")

        with open(file_path, 'r', encoding='utf-8') as f:
            instance = json.load(f)

        # 验证必要字段
        required_keys = ['problem_type', 'parameters', 'processing_times']
        for key in required_keys:
            if key not in instance:
                raise ValueError(f"算例格式错误，缺少必要字段: {key}")

        # 兼容旧算例（充气时间二维化）
        if len(np.array(instance['processing_times']['inflation_time']).shape) == 1:
            num_workpieces = instance['parameters']['num_workpieces']
            num_components = instance['parameters']['num_components']
            old_inflation = np.array(instance['processing_times']['inflation_time'])
            new_inflation = np.tile(old_inflation[:, np.newaxis], (1, num_components))
            instance['processing_times']['inflation_time'] = new_inflation.tolist()
            print(f"警告：旧算例已自动转换为多组分充装格式")

        # 补充缺失参数
        if 'num_clean_eq' not in instance['parameters']:
            instance['parameters']['num_clean_eq'] = 1
        if 'num_shake_eq' not in instance['parameters']:
            instance['parameters']['num_shake_eq'] = 1
        if 'clean_time' not in instance['processing_times']:
            instance['processing_times']['clean_time'] = 5
        if 'shake_unit_time' not in instance['processing_times']:
            instance['processing_times']['shake_unit_time'] = 2

        return instance

    @staticmethod
    def list_instances(directory: str = ".") -> List[str]:
        instance_files = []
        for root, _, files in os.walk(directory):
            for file in files:
                if file.endswith(".json"):
                    try:
                        with open(os.path.join(root, file), 'r') as f:
                            data = json.load(f)
                            if data.get('problem_type') == 'gas_scheduling':
                                instance_files.append(os.path.join(root, file))
                    except:
                        continue
        return sorted(instance_files)


class BaseSchedulingAlgorithm:
    """调度算法基类"""

    def __init__(self,
                 instance: Dict,
                 max_iter: int = 100,
                 name: str = "算法"):
        params = instance['parameters']
        processing_times = instance['processing_times']

        # 基础参数
        self.num_workpieces = params['num_workpieces']
        self.num_components = params['num_components']
        self.num_inflation_eq = params['num_inflation_eq']
        self.num_analysis_eq_per_component = params['num_analysis_eq_per_component']
        self.num_clean_eq = params['num_clean_eq']
        self.num_shake_eq = params['num_shake_eq']

        # 加工时间
        self.inflation_time = np.array(processing_times['inflation_time'])
        self.analysis_time = np.array(processing_times['analysis_time'])
        self.clean_time = processing_times['clean_time']
        self.shake_unit_time = processing_times['shake_unit_time']
        self.shake_time = np.array([self.num_components * self.shake_unit_time
                                    for _ in range(self.num_workpieces)])

        self.max_iter = max_iter
        self.name = name
        self.best_solution = None
        self.best_time = float('inf')
        self.iteration_best_times = []
        self.execution_time = 0

    def initialize_solution(self) -> Dict:
        """初始化解（多组分充装适配）"""
        # 1. 清洗阶段
        clean_order = random.sample(range(self.num_workpieces), self.num_workpieces)
        clean_eq = [random.randint(0, self.num_clean_eq - 1)
                    for _ in range(self.num_workpieces)]

        # 2. 充装阶段
        inflation_order = []  # (workpiece, component) 元组
        for comp in range(self.num_components):
            comp_order = [(wp, comp) for wp in range(self.num_workpieces)]
            random.shuffle(comp_order)
            inflation_order.extend(comp_order)
        inflation_eq = [
            [random.randint(0, self.num_inflation_eq - 1)
             for _ in range(self.num_components)]
            for _ in range(self.num_workpieces)
        ]

        # 3. 摇匀阶段
        shake_order = random.sample(range(self.num_workpieces), self.num_workpieces)
        shake_eq = [random.randint(0, self.num_shake_eq - 1)
                    for _ in range(self.num_workpieces)]

        # 4. 分析阶段
        analysis_order = random.sample(range(self.num_workpieces), self.num_workpieces)
        component_groups = []
        group_test_order = []
        analysis_eq = []
        for _ in range(self.num_workpieces):
            max_groups = random.randint(1, self.num_components)
            groups = [random.randint(0, max_groups - 1) for _ in range(self.num_components)]
            component_groups.append(groups)
            unique_groups = list(sorted(set(groups)))
            random.shuffle(unique_groups)
            group_test_order.append(unique_groups)
            eq_for_workpiece = []
            for group_id in unique_groups:
                components_in_group = [j for j, g in enumerate(groups) if g == group_id]
                eq_for_group = [random.randint(0, self.num_analysis_eq_per_component - 1)
                                for _ in components_in_group]
                eq_for_workpiece.append(eq_for_group)
            analysis_eq.append(eq_for_workpiece)

        return {
            'clean_order': clean_order,
            'clean_eq': clean_eq,
            'inflation_order': inflation_order,
            'inflation_eq': inflation_eq,
            'shake_order': shake_order,
            'shake_eq': shake_eq,
            'analysis_order': analysis_order,
            'component_groups': component_groups,
            'group_test_order': group_test_order,
            'analysis_eq': analysis_eq
        }

    def calculate_completion_time(self, solution: Dict) -> float:
        """计算四阶段总时间"""
        # 1. 清洗阶段
        clean_machine_time = [0] * self.num_clean_eq
        workpiece_clean_end = {}
        for workpiece in solution['clean_order']:
            eq = solution['clean_eq'][workpiece]
            start_time = clean_machine_time[eq]
            end_time = start_time + self.clean_time
            clean_machine_time[eq] = end_time
            workpiece_clean_end[workpiece] = end_time
        clean_completion_time = max(clean_machine_time)

        # 2. 充装阶段
        inflation_machine_time = [0] * self.num_inflation_eq
        workpiece_comp_inflate_end = {
            wp: {comp: 0 for comp in range(self.num_components)}
            for wp in range(self.num_workpieces)
        }
        for (wp, comp) in solution['inflation_order']:
            eq = solution['inflation_eq'][wp][comp]
            start_time = max(
                inflation_machine_time[eq],
                workpiece_clean_end.get(wp, 0),
                max(workpiece_comp_inflate_end[wp].values())
            )
            end_time = start_time + self.inflation_time[wp][comp]
            inflation_machine_time[eq] = end_time
            workpiece_comp_inflate_end[wp][comp] = end_time
        inflation_completion_time = max(inflation_machine_time)
        workpiece_inflate_end = {
            wp: max(workpiece_comp_inflate_end[wp].values())
            for wp in range(self.num_workpieces)
        }

        # 3. 摇匀阶段
        shake_machine_time = [0] * self.num_shake_eq
        workpiece_shake_end = {}
        for workpiece in solution['shake_order']:
            eq = solution['shake_eq'][workpiece]
            start_time = max(shake_machine_time[eq], workpiece_inflate_end.get(workpiece, 0))
            end_time = start_time + self.shake_time[workpiece]
            shake_machine_time[eq] = end_time
            workpiece_shake_end[workpiece] = end_time
        shake_completion_time = max(shake_machine_time)

        # 4. 分析阶段
        analysis_machine_time = [[0 for _ in range(self.num_analysis_eq_per_component)]
                                 for _ in range(self.num_components)]
        for workpiece in solution['analysis_order']:
            if workpiece < 0 or workpiece >= self.num_workpieces:
                continue
            start_analysis_time = workpiece_shake_end.get(workpiece, 0)
            group_test_order = solution['group_test_order'][workpiece]
            analysis_eq = solution['analysis_eq'][workpiece]

            for group_idx, group_id in enumerate(group_test_order):
                if group_idx < 0 or group_idx >= len(analysis_eq):
                    continue
                components_in_group = [j for j, g in enumerate(solution['component_groups'][workpiece])
                                       if g == group_id]
                eq_for_group = analysis_eq[group_idx]
                if len(eq_for_group) != len(components_in_group):
                    eq_for_group = [random.randint(0, self.num_analysis_eq_per_component - 1)
                                    for _ in components_in_group]
                    analysis_eq[group_idx] = eq_for_group

                group_ready_time = start_analysis_time
                for comp_idx, component in enumerate(components_in_group):
                    if comp_idx < 0 or comp_idx >= len(eq_for_group):
                        continue
                    eq = eq_for_group[comp_idx]
                    group_ready_time = max(group_ready_time, analysis_machine_time[component][eq])

                max_analysis_time = max(
                    self.analysis_time[workpiece][component]
                    for comp_idx, component in enumerate(components_in_group)
                    if comp_idx < len(eq_for_group)
                ) if components_in_group else 0
                group_end_time = group_ready_time + max_analysis_time

                for comp_idx, component in enumerate(components_in_group):
                    if comp_idx < 0 or comp_idx >= len(eq_for_group):
                        continue
                    eq = eq_for_group[comp_idx]
                    analysis_machine_time[component][eq] = group_end_time
                start_analysis_time = group_end_time

        analysis_completion_time = max(max(machine_times) for machine_times in analysis_machine_time)

        return max(clean_completion_time, inflation_completion_time,
                   shake_completion_time, analysis_completion_time)

    def get_schedule_details(self, solution: Dict) -> Tuple[Dict, Dict, Dict, Dict, float]:
        """返回四阶段调度详情及分析阶段偏移量"""
        # 1. 清洗阶段
        clean_details = {eq: [] for eq in range(self.num_clean_eq)}
        clean_machine_time = [0] * self.num_clean_eq
        workpiece_clean_end = {}
        for workpiece in solution['clean_order']:
            eq = solution['clean_eq'][workpiece]
            start = clean_machine_time[eq]
            end = start + self.clean_time
            clean_machine_time[eq] = end
            workpiece_clean_end[workpiece] = end
            clean_details[eq].append({'workpiece': workpiece, 'start': start, 'end': end})

        # 2. 充装阶段
        inflation_details = {eq: [] for eq in range(self.num_inflation_eq)}
        inflation_machine_time = [0] * self.num_inflation_eq
        workpiece_comp_inflate_end = {
            wp: {comp: 0 for comp in range(self.num_components)}
            for wp in range(self.num_workpieces)
        }
        for (wp, comp) in solution['inflation_order']:
            eq = solution['inflation_eq'][wp][comp]
            start = max(
                inflation_machine_time[eq],
                workpiece_clean_end.get(wp, 0),
                max(workpiece_comp_inflate_end[wp].values())
            )
            end = start + self.inflation_time[wp][comp]
            inflation_machine_time[eq] = end
            workpiece_comp_inflate_end[wp][comp] = end
            inflation_details[eq].append({
                'workpiece': wp, 'component': comp, 'start': start, 'end': end,
                'duration': self.inflation_time[wp][comp]
            })
        workpiece_inflate_end = {
            wp: max(workpiece_comp_inflate_end[wp].values())
            for wp in range(self.num_workpieces)
        }

        # 3. 摇匀阶段
        shake_details = {eq: [] for eq in range(self.num_shake_eq)}
        shake_machine_time = [0] * self.num_shake_eq
        workpiece_shake_end = {}
        for workpiece in solution['shake_order']:
            eq = solution['shake_eq'][workpiece]
            start = max(shake_machine_time[eq], workpiece_inflate_end.get(workpiece, 0))
            end = start + self.shake_time[workpiece]
            shake_machine_time[eq] = end
            workpiece_shake_end[workpiece] = end
            shake_details[eq].append({
                'workpiece': workpiece, 'start': start, 'end': end,
                'group_count': self.num_components
            })

        # 4. 分析阶段
        analysis_details = {c: {eq: [] for eq in range(self.num_analysis_eq_per_component)}
                            for c in range(self.num_components)}
        analysis_machine_time = [[0] * self.num_analysis_eq_per_component for _ in range(self.num_components)]
        all_analysis_starts = []

        for workpiece in solution['analysis_order']:
            if workpiece < 0 or workpiece >= self.num_workpieces:
                continue
            start_analysis_time = workpiece_shake_end.get(workpiece, 0)
            group_test_order = solution['group_test_order'][workpiece]
            analysis_eq = solution['analysis_eq'][workpiece]

            for group_idx, group_id in enumerate(group_test_order):
                if group_idx < 0 or group_idx >= len(analysis_eq):
                    continue
                components_in_group = [j for j, g in enumerate(solution['component_groups'][workpiece]) if
                                       g == group_id]
                eq_for_group = analysis_eq[group_idx]
                if len(eq_for_group) != len(components_in_group):
                    eq_for_group = [random.randint(0, self.num_analysis_eq_per_component - 1) for _ in
                                    components_in_group]

                group_ready_time = start_analysis_time
                for comp_idx, component in enumerate(components_in_group):
                    if comp_idx < 0 or comp_idx >= len(eq_for_group):
                        continue
                    eq = eq_for_group[comp_idx]
                    group_ready_time = max(group_ready_time, analysis_machine_time[component][eq])

                all_analysis_starts.append(group_ready_time)

                max_analysis_time = max(
                    self.analysis_time[workpiece][component]
                    for comp_idx, component in enumerate(components_in_group)
                    if comp_idx < len(eq_for_group)
                ) if components_in_group else 0
                group_end_time = group_ready_time + max_analysis_time

                for comp_idx, component in enumerate(components_in_group):
                    if comp_idx < 0 or comp_idx >= len(eq_for_group):
                        continue
                    eq = eq_for_group[comp_idx]
                    analysis_machine_time[component][eq] = group_end_time
                    analysis_details[component][eq].append({
                        'workpiece': workpiece, 'start': group_ready_time, 'end': group_end_time,
                        'group': group_id, 'individual_time': self.analysis_time[workpiece][component]
                    })
                start_analysis_time = group_end_time

        analysis_start_offset = min(all_analysis_starts) if all_analysis_starts else 0
        return clean_details, inflation_details, shake_details, analysis_details, analysis_start_offset

    def plot_gantt_chart(self, solution: Dict = None, title_suffix: str = ""):
        """绘制甘特图"""
        if solution is None:
            solution = self.best_solution
        if solution is None:
            print(f"{self.name} 没有可用的调度方案进行绘制")
            return

        clean_details, inflation_details, shake_details, analysis_details, analysis_start_offset = self.get_schedule_details(
            solution)
        total_global_time = self.calculate_completion_time(solution)

        # 颜色分配
        base_colors = plt.cm.tab10(np.linspace(0, 1, min(10, self.num_workpieces)))
        wp_comp_colors = {}
        for wp in range(self.num_workpieces):
            base_color = base_colors[wp % 10]
            for comp in range(self.num_components):
                lightness = 0.8 + (comp % 3) * 0.15  # 轻微调整同工件不同组分的亮度
                wp_comp_colors[(wp, comp)] = (
                    min(base_color[0] * lightness, 1.0),
                    min(base_color[1] * lightness, 1.0),
                    min(base_color[2] * lightness, 1.0)
                )

        # 创建子图
        top_fig_height = (self.num_clean_eq + self.num_inflation_eq + self.num_shake_eq) * 0.8 + 2
        analysis_eq_total = self.num_components * self.num_analysis_eq_per_component
        bottom_fig_height = analysis_eq_total * 0.6 + 1
        fig, (ax_top, ax_bottom) = plt.subplots(
            2, 1, figsize=(16, top_fig_height + bottom_fig_height + 1),
            gridspec_kw={'height_ratios': [top_fig_height, bottom_fig_height]}
        )
        fig.suptitle(f'{self.name} 气体调度甘特图 {title_suffix}', fontsize=16, y=0.98)

        # 上方子图：清洗+充装+摇匀（全局时间）
        y_base = 0
        stage_y_ranges = {}

        # 清洗阶段
        stage_name = "清洗阶段"
        stage_eq_count = self.num_clean_eq
        stage_y_start = y_base + 1
        stage_y_end = stage_y_start + stage_eq_count
        stage_y_ranges[stage_name] = (stage_y_start, stage_y_end)

        for eq_idx, eq in enumerate(sorted(clean_details.keys())):
            y_pos = stage_y_start + eq_idx
            ax_top.text(-total_global_time * 0.02, y_pos, f'{stage_name} 设备{eq}',
                        ha='right', va='center', fontweight='bold', fontsize=9)
            for task in clean_details[eq]:
                wp = task['workpiece']
                ax_top.barh(
                    y_pos, task['end'] - task['start'], left=task['start'], height=0.6,
                    color=base_colors[wp % 10], edgecolor='black', alpha=0.9
                )
                ax_top.text(
                    task['start'] + (task['end'] - task['start']) / 2, y_pos,
                    f'瓶{wp}', ha='center', va='center',
                    color='white', fontweight='bold', fontsize=8
                )
        y_base = stage_y_end + 1  # 预留阶段间隔

        # 充装阶段
        stage_name = "充装阶段"
        stage_eq_count = self.num_inflation_eq
        stage_y_start = y_base
        stage_y_end = stage_y_start + stage_eq_count
        stage_y_ranges[stage_name] = (stage_y_start, stage_y_end)

        for eq_idx, eq in enumerate(sorted(inflation_details.keys())):
            y_pos = stage_y_start + eq_idx
            ax_top.text(-total_global_time * 0.02, y_pos, f'{stage_name} 设备{eq}',
                        ha='right', va='center', fontweight='bold', fontsize=9)
            for task in inflation_details[eq]:
                wp = task['workpiece']
                comp = task['component']
                ax_top.barh(
                    y_pos, task['end'] - task['start'], left=task['start'], height=0.6,
                    color=wp_comp_colors[(wp, comp)], edgecolor='black', alpha=0.9
                )
                ax_top.text(
                    task['start'] + (task['end'] - task['start']) / 2, y_pos,
                    f'瓶{wp}-组{comp}', ha='center', va='center',
                    color='white', fontweight='bold', fontsize=7
                )
        y_base = stage_y_end + 1

        # 摇匀阶段
        stage_name = "摇匀阶段"
        stage_eq_count = self.num_shake_eq
        stage_y_start = y_base
        stage_y_end = stage_y_start + stage_eq_count
        stage_y_ranges[stage_name] = (stage_y_start, stage_y_end)

        for eq_idx, eq in enumerate(sorted(shake_details.keys())):
            y_pos = stage_y_start + eq_idx
            ax_top.text(-total_global_time * 0.02, y_pos, f'{stage_name} 设备{eq}',
                        ha='right', va='center', fontweight='bold', fontsize=9)
            for task in shake_details[eq]:
                wp = task['workpiece']
                ax_top.barh(
                    y_pos, task['end'] - task['start'], left=task['start'], height=0.6,
                    color=base_colors[wp % 10], edgecolor='black', alpha=0.9
                )
                ax_top.text(
                    task['start'] + (task['end'] - task['start']) / 2, y_pos,
                    f'瓶{wp}', ha='center', va='center',
                    color='white', fontweight='bold', fontsize=8
                )

        # 上方子图设置
        ax_top.set_xlim(0, total_global_time * 1.05)
        ax_top.set_ylim(0, stage_y_end + 1)
        ax_top.set_xlabel('时间')
        ax_top.set_yticks([])
        ax_top.grid(axis='x', linestyle='--', alpha=0.7)

        # 下方子图：分析阶段（相对时间，便于查看细节）
        y_base = 0
        for component in sorted(analysis_details.keys()):
            comp_details = analysis_details[component]
            for eq in sorted(comp_details.keys()):
                y_pos = y_base + 0.5
                ax_bottom.text(-total_global_time * 0.01, y_pos, f'组分{component} 设备{eq}',
                               ha='right', va='center', fontweight='bold', fontsize=8)

                for task in comp_details[eq]:
                    wp = task['workpiece']
                    # 转换为相对时间（减去分析阶段起始偏移，便于观察）
                    start_rel = task['start'] - analysis_start_offset
                    end_rel = task['end'] - analysis_start_offset
                    ax_bottom.barh(
                        y_pos, end_rel - start_rel, left=start_rel, height=0.4,
                        color=base_colors[wp % 10], edgecolor='black', alpha=0.9
                    )
                    ax_bottom.text(
                        start_rel + (end_rel - start_rel) / 2, y_pos,
                        f'瓶{wp}-组{task["group"]}', ha='center', va='center',
                        color='white', fontweight='bold', fontsize=7
                    )
                y_base += 1

        # 下方子图设置
        analysis_max_time = total_global_time - analysis_start_offset if analysis_start_offset else total_global_time
        ax_bottom.set_xlim(0, analysis_max_time * 1.05)
        ax_bottom.set_ylim(0, y_base + 1)
        ax_bottom.set_xlabel(f'相对时间（相对于分析阶段起始点：{analysis_start_offset:.1f}）')
        ax_bottom.set_yticks([])
        ax_bottom.grid(axis='x', linestyle='--', alpha=0.7)

        # 添加图例
        legend_patches = []
        for wp in range(self.num_workpieces):
            legend_patches.append(mpatches.Patch(
                color=base_colors[wp % 10], label=f'工件 {wp}'
            ))
        plt.figlegend(handles=legend_patches, loc='lower center', ncol=min(10, self.num_workpieces),
                      bbox_to_anchor=(0.5, 0.01), bbox_transform=plt.gcf().transFigure)

        plt.tight_layout()
        plt.subplots_adjust(bottom=0.15)  # 预留图例空间
        plt.show()